package cn.pugle.oj.leetcode;

import cn.pugle.oj.catalog.TreeProblem;

import java.util.HashMap;
import java.util.Map;

/**
 * https://leetcode.com/problems/implement-trie-prefix-tree/
 *
 * @author tzp
 * @since 2020/10/14
 */
public class LC208 implements TreeProblem {
    public static class TrieNode {
        public boolean isWordEnd;
        public Map<Character, TrieNode> children;

        public TrieNode() {
            this.isWordEnd = false;
            this.children = new HashMap<>();
        }
    }

    public static class Trie {
        TrieNode root;//要有个dummyRoot

        /**
         * Initialize your data structure here.
         */
        public Trie() {
            root = new TrieNode();
        }

        /**
         * Inserts a wd into the trie.
         */
        public void insert(String word) {
            TrieNode cur = root;
            for (int i = 0; i < word.length(); i++) {
                cur.children.computeIfAbsent(word.charAt(i), character -> new TrieNode());
                cur = cur.children.get(word.charAt(i));
            }
            cur.isWordEnd = true;
        }

        /**
         * Returns if the word is in the trie.
         */
        public boolean search(String word) {
            TrieNode cur = root;
            for (int i = 0; i < word.length(); i++) {
                if (!cur.children.containsKey(word.charAt(i))) {
                    return false;
                } else {
                    cur = cur.children.get(word.charAt(i));
                }
            }
            return cur.isWordEnd;
        }

        /**
         * Returns if there is any word in the trie that starts with the given prefix.
         */
        public boolean startsWith(String prefix) {
            TrieNode cur = root;
            for (int i = 0; i < prefix.length(); i++) {
                if (!cur.children.containsKey(prefix.charAt(i))) {
                    return false;
                } else {
                    cur = cur.children.get(prefix.charAt(i));
                }
            }
            return cur != null;
        }
    }

    public static void main(String[] args) {
        Trie obj = new Trie();
        obj.insert("rome");
        obj.insert("romantic");
        obj.insert("robe");
        System.out.println(obj.search("ro"));
        System.out.println(obj.search("robe"));
        System.out.println(obj.startsWith("rob"));
        System.out.println(obj.startsWith("ra"));
    }
}
